# 选择gpu运行
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 训练次数
EPOCHS = 5
# 模型实例化
model = ResNet(ResidualBlock,num_classes= 4 * 9).to(DEVICE)
# 优化器：更新模型参数,使训练结果达到最优值
optimizer = optim.Adam(model.parameters())
total_loss = []
# 这一步借用了tqdm实现了进度条打印的功能
dataloader = tqdm(train_loader, total=len(train_loader))
# 定义损失函数
loss_function = torch.nn.CrossEntropyLoss()
# 启动训练
model.train()
for epoch in range(EPOCHS):
    for (image, label) in dataloader:
        image = image.to(DEVICE) # 部署到DEVICE
        label = label.to(DEVICE)
        # 向前传播
        output = model(image) #(8,36)

        # 9是分类, 4*9就是4个9分类的结果（要分类才需要这些操作）
        # 将模型的输出 output 和标签 label 进行维度的变换,
        # 使它们能够被送入交叉熵损失函数进行计算。
        output = output.view(BATCH_SIZE * 4, 9)
        label = label.view(-1)
        # 计算损失
        loss = loss_function(output, label)
        total_loss.append(loss.item())
        
        optimizer.zero_grad()# 梯度初始化为0        
        loss.backward()# 反向传播        
        optimizer.step()# 优化器更新
    # 保存模型
    torch.save(model.state_dict(), './model.pkl')
    # 保存优化器
    torch.save(optimizer.state_dict(), './optimizer.pkl')
    mean_loss =np.mean(total_loss)
    print(f"第{epoch + 1}次epoch---损失: {mean_loss}") 
